import uuid
import datetime
import json
import threading
from typing import List

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect, Depends, Cookie
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles

from pydantic import BaseModel
from starlette.concurrency import run_in_threadpool

from ai.ai_util import AIClient

app = FastAPI()

# 配置模板目录
templates = Jinja2Templates(directory="templates")

# 将 static 目录映射到 /static 路径
app.mount("/static", StaticFiles(directory="static"), name="static")

# 创建 AIClient 实例
ai_client = AIClient()

# 全局会话历史字典和锁（用于存储每个用户的聊天历史）
session_histories = {}
history_lock = threading.Lock()

# 如果请求中没有 session_id cookie，则生成一个新的
def get_session_id(request: Request, session_id: str = Cookie(None)):
    if not session_id:
        session_id = str(uuid.uuid4())
    return session_id

# 定义消息请求体（HTTP接口使用）
class Message(BaseModel):
    text: str

# HTTP 接口：普通聊天请求（非流式实时）
@app.post("/chat")
async def chat(message: Message, request: Request, session_id: str = Depends(get_session_id)):
    user_message = message.text
    print(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Received message from session {session_id}: {user_message}")

    # 获取当前用户的聊天历史
    with history_lock:
        history = session_histories.get(session_id, [])
    # 调用 AIClient 的同步接口放入线程池中，避免阻塞事件循环
    try:
        ai_response, new_history = await run_in_threadpool(ai_client.get_ai_response, user_message, history)
    except Exception as e:
        return {"error": f"Failed to get AI response: {str(e)}"}
    # 更新聊天历史
    with history_lock:
        session_histories[session_id] = new_history
    print(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Generated AI response for session {session_id}")
    return {"response": ai_response}

# WebSocket 接口：实时流式聊天
@app.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    # 接受 WebSocket 连接
    await websocket.accept()

    # 尝试从查询参数中获取 session_id，否则生成新的
    session_id = websocket.query_params.get("session_id")
    if not session_id:
        session_id = str(uuid.uuid4())
    print(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] WebSocket connected, session: {session_id}")

    # 初始化该 session 的历史（如果不存在）
    with history_lock:
        if session_id not in session_histories:
            session_histories[session_id] = []

    try:
        while True:
            # 接收用户消息
            user_message = await websocket.receive_text()
            print(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] WS received from {session_id}: {user_message}")

            # 获取当前会话历史
            with history_lock:
                history = session_histories.get(session_id, [])
            # 调用 AI 服务，要求流式返回（同步调用放入线程池）
            response = await run_in_threadpool(ai_client.post_http_request, user_message, history, 2048, 0.95, 1, 0.8, False, True)

            # 逐步读取流式响应，并发送给客户端
            for output, new_history in ai_client.get_streaming_response(response):
                # 更新历史（每个 chunk 都可更新，也可以只在整次会话结束后更新）
                with history_lock:
                    session_histories[session_id] = new_history
                print(f">>>{output}")
                await websocket.send_text(output)
            # 标记一次回复结束（可选）
            await websocket.send_text("[END]")
    except WebSocketDisconnect:
        print(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] WebSocket disconnected, session: {session_id}")

# 主页，返回聊天页面模板
@app.get("/", response_class=HTMLResponse)
async def get_chat_page(request: Request):
    return templates.TemplateResponse("chat.html", {"request": request})
